// Copyright (C) joyinter
// All rights reserved.
// 
// Author : Jiang Li
// Date   : 2005-5-12

#include "stdafx.h"
#include "modcheck.h"
#include "apihelp.h"
#include "crc.h"
#include "strcrypt.h"

namespace JoyGuard
{

//##ModelId=44CF17BE0324
ModCheck::ModCheck()
{
    m_process   = NULL;
    m_processId = 0;
}

//##ModelId=44CF17BE032C
ModCheck::~ModCheck()
{
    ModDataMap::iterator i = m_datas.begin();
    ModDataMap::iterator e = m_datas.end();
    for(; i != e; ++ i)
    {
        delete i->second;
    }
    m_datas.clear();
}

//##ModelId=44CF17BE032F
void ModCheck::LoadCheck(BYTE* check, int checkLen)
{
    BYTE* encBuf = check;
    StrCrypt::DecryptData(encBuf, encBuf, checkLen);

    DWORD len = 0;
    ModCheck::ModData* modData = new ModCheck::ModData();
    modData->m_headerHash = *(DWORD*)check; check += sizeof(DWORD);

    len = *(DWORD*)check; check += sizeof(DWORD);
    modData->m_headerData.insert(modData->m_headerData.end(), check, check+len); check += len;

    len = *(DWORD*)check; check += sizeof(DWORD);
    modData->m_codeDatas.insert(modData->m_codeDatas.end(),
        (ModCheck::CodeData*)check, (ModCheck::CodeData*)check+len);

    m_datas.insert(std::pair<DWORD, ModCheck::ModData*>(modData->m_headerHash, modData));

    StrCrypt::EncryptData(encBuf, encBuf, checkLen);
}

//##ModelId=44CF17BE031F
void ModCheck::ReadMemory(void* dst, void* src, int len)
{
    __super::ReadMemory(m_process, (BYTE*)dst, (BYTE*)src, len);
}

//##ModelId=44CF17BE0311
void ModCheck::OnMatch(ModData*)
{
    m_error = EC_CHECK_IMAGE;
#ifdef _DEBUG
//#if 0
    MessageBox(NULL, "find game hooks.", "", MB_OK);
    throw ECheckStop();
#else
    HANDLE hModuleSnap = INVALID_HANDLE_VALUE;
    hModuleSnap = ApiHelp::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, m_processId);
    if(hModuleSnap == INVALID_HANDLE_VALUE)
    {
        return;
    }
    MODULEENTRY32 me32;
    me32.dwSize = sizeof( MODULEENTRY32 );

    if(!ApiHelp::Module32First(hModuleSnap, &me32))
    {
        ::CloseHandle(hModuleSnap);
        return;
    }

    BYTE* modBase = NULL;

    do
    {
        const char* p = me32.szModule+strlen(me32.szModule)-4;
        if(p < me32.szModule)
        {
            continue;
        }
        if( p[0] != '.' || 
            p[1] != 'e' ||
            p[2] != 'x' ||
            p[3] != 'e' )
        {
            continue;
        }
        modBase = me32.modBaseAddr;
        break;
    }
    while(ApiHelp::Module32Next(hModuleSnap, &me32));

    ::CloseHandle(hModuleSnap);

    if(modBase == NULL)
    {
        return;
    }

    IMAGE_DOS_HEADER dosHeader;
    ReadMemory(&dosHeader, modBase, sizeof(dosHeader));

    IMAGE_NT_HEADERS ntHeaders;
    ReadMemory(&ntHeaders, modBase+dosHeader.e_lfanew, sizeof(ntHeaders));

    int sectsSize = ntHeaders.FileHeader.NumberOfSections*sizeof(IMAGE_SECTION_HEADER);
    std::vector<BYTE> headerBuf(sizeof(IMAGE_NT_HEADERS) + sectsSize);

    memcpy(&headerBuf[0], &ntHeaders, sizeof(IMAGE_NT_HEADERS));
    
    ReadMemory(&headerBuf[sizeof(IMAGE_NT_HEADERS)], 
        modBase+dosHeader.e_lfanew+sizeof(IMAGE_NT_HEADERS), sectsSize);

    // sectName is ".text"
    char sectName[8];
    sectName[1] = 't';
    sectName[4] = 's';
    IMAGE_NT_HEADERS* pntHeaders = (IMAGE_NT_HEADERS*)&headerBuf[0];
    sectName[0] = '.';
    sectName[2] = 'f';

    // find a section match the sectName
    int sectCount = pntHeaders->FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER* sectheaders = (IMAGE_SECTION_HEADER*)(pntHeaders+1);
    IMAGE_SECTION_HEADER* foundHeader = NULL;
    sectName[3] = 'x';
    sectName[5] = 0;
    ++ sectName[4];
    -- sectName[2];
    for(int i = 0; i < (int)sectCount; i ++)
    {
        IMAGE_SECTION_HEADER* header = sectheaders+i;
        if(header->Characteristics == 0)
        {
            return;
        }
        if(strcmp((char*)header->Name, sectName) == 0)
        {
            foundHeader = header;
            break;
        }
    }
    if(foundHeader == NULL)
    {
        return;
    }

    BYTE* startAddr = modBase + foundHeader->VirtualAddress;
    BYTE* stopAddr = startAddr + foundHeader->Misc.VirtualSize;

    DWORD oldPro = 0;
    if(!ApiHelp::VirtualProtectEx(m_process, startAddr, stopAddr-startAddr, PAGE_EXECUTE_READWRITE, &oldPro))
    {
        return;
    }

    BYTE wbuf[4096];
    memset(wbuf, 0xCCCCCCCC, sizeof(wbuf));
    for(; startAddr < stopAddr; startAddr += sizeof(wbuf))
    {
        BYTE* e = startAddr + sizeof(wbuf);
        if(e > stopAddr)
        {
            e = stopAddr;
        }
        SIZE_T written = 0;
        if(!ApiHelp::WriteProcessMemory(m_process, startAddr, wbuf, e-startAddr, &written))
        {
            //::GetTickCount();
        }
    }
#endif
}

//##ModelId=44CF17BE0314
void ModCheck::OnCheck(void* pme32)
{
    MODULEENTRY32& me32 = *(MODULEENTRY32*)pme32;

    IMAGE_DOS_HEADER dosHeader;
    ReadMemory(&dosHeader, me32.modBaseAddr, sizeof(dosHeader));

    IMAGE_NT_HEADERS ntHeaders;
    ReadMemory(&ntHeaders, me32.modBaseAddr+dosHeader.e_lfanew, sizeof(ntHeaders));

    int sectsSize = ntHeaders.FileHeader.NumberOfSections*sizeof(IMAGE_SECTION_HEADER);
    std::vector<BYTE> headerBuf(sizeof(IMAGE_NT_HEADERS) + sectsSize);

    memcpy(&headerBuf[0], &ntHeaders, sizeof(IMAGE_NT_HEADERS));
    
    ReadMemory(&headerBuf[sizeof(IMAGE_NT_HEADERS)], 
        me32.modBaseAddr+dosHeader.e_lfanew+sizeof(IMAGE_NT_HEADERS), sectsSize);

    Crc headerSum;
    headerSum.Update(&headerBuf[0], (int)headerBuf.size());
    DWORD headerHash = headerSum.Final();

    std::vector<BYTE> codeBuf;
    ModDataMap::iterator it = m_datas.find(headerHash);
    ModDataMap::iterator e = m_datas.end();
    for( ; it != e && it->first == headerHash; ++ it)
    {
        ModData* d = it->second;
        int s = (int)headerBuf.size();
        if(s != (int)d->m_headerData.size())
        {
            continue;
        }
        if(memcmp(&headerBuf[0], &d->m_headerData[0], s) != 0)
        {
            continue;
        }

        // check code
        if(!d->m_codeDatas.empty())
        {
            bool match = true;
            for(int i = 0; i < (int)d->m_codeDatas.size(); ++ i)
            {
                CodeData& cd = d->m_codeDatas[i];
                codeBuf.resize(cd.m_size);
                ReadMemory(&codeBuf[0], me32.modBaseAddr+cd.m_rva, (int)cd.m_size);
                Crc codeSum;
                codeSum.Update(&codeBuf[0], (int)cd.m_size);
                DWORD h = codeSum.Final();
                if(h != cd.m_hash)
                {
                    match = false;
                    break;
                }
            }
            if(!match)
            {
                continue;
            }
        }

        // find matched module
        OnMatch(d);
        break;
    }
}

//##ModelId=44CF17BE031D
void ModCheck::CheckProc()
{
    while(!m_stopped)
    {
        HANDLE hModuleSnap = INVALID_HANDLE_VALUE;
        MODULEENTRY32 me32;
        hModuleSnap = ApiHelp::CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, m_processId);
        if(hModuleSnap != INVALID_HANDLE_VALUE)
        {
            me32.dwSize = sizeof( MODULEENTRY32 );

            if(ApiHelp::Module32First(hModuleSnap, &me32))
            {
                do
                {
                    try
                    {
                        if(m_checks.find(me32.modBaseAddr) == m_checks.end())
                        {
                            m_checks.insert(me32.modBaseAddr);
                            OnCheck(&me32);
                        }
                    }
                    catch(...)
                    {
                        if(m_stopped)
                        {
                            break;
                        }
                        else if(m_error == EC_CHECK_IMAGE)
                        {
                            // FIX ME
                            // do hook clear here.
                            break;
                        }
                    }
                }
                while(ApiHelp::Module32Next(hModuleSnap, &me32));
            }

            ::CloseHandle(hModuleSnap);
        }
        
        ApiHelp::WaitForSingleObject(m_waitEvent, m_checkDelay);
    }
}

//##ModelId=44CF17BE032E
bool ModCheck::Check()
{
    return true;
}

//##ModelId=44CF17BE0332
bool ModCheck::BuildModDataFromFile(const char* fileName, ModData& modData)
{
    // read data from pe file
    HANDLE peFile = ::CreateFile(fileName, GENERIC_READ,
        FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if(peFile == INVALID_HANDLE_VALUE)
    {
        return false;
    }

    DWORD dwSize = ::GetFileSize(peFile, NULL);
    if(dwSize < sizeof(IMAGE_DOS_HEADER))
    {
        return false;
    }
    DWORD dwReaded = 0;
    std::vector<BYTE> peBuf(dwSize);
    if(!::ReadFile(peFile, &peBuf[0], dwSize, &dwReaded, NULL))
    {
        return false;
    }

    BYTE* baseAddr = &peBuf[0];

    // get dos header
    IMAGE_DOS_HEADER* dosHeader    = (IMAGE_DOS_HEADER*)baseAddr;
    if(dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
    {
        return NULL;
    }
    if(dosHeader->e_lfanew <= 0)
    {
        return false;
    }

    // get nt headers
    IMAGE_NT_HEADERS* ntHeaders   = (IMAGE_NT_HEADERS*)(baseAddr+dosHeader->e_lfanew);
    if(ntHeaders->Signature != IMAGE_NT_SIGNATURE)
    {
        return false;
    }

    IMAGE_SECTION_HEADER* sectHeader = (IMAGE_SECTION_HEADER*)(ntHeaders+1);

    std::vector<BYTE> imgBuf;
    imgBuf.resize(ntHeaders->OptionalHeader.SizeOfImage);
    BYTE* headerEnd = (BYTE*)(sectHeader+ntHeaders->FileHeader.NumberOfSections);
    std::copy(baseAddr, headerEnd, &imgBuf[0]);

    for(; sectHeader < (IMAGE_SECTION_HEADER*)headerEnd; ++ sectHeader)
    {
        memcpy(&imgBuf[sectHeader->VirtualAddress],
            baseAddr+sectHeader->PointerToRawData,
            sectHeader->Misc.VirtualSize);
    }

    return BuildModDataFromMemory(&imgBuf[0], modData);
}

//##ModelId=44CF17BE033F
bool ModCheck::BuildModDataFromModule(HANDLE hProcess, BYTE* pMod, ModData& modData)
{
    SIZE_T readed;
    IMAGE_DOS_HEADER dosHeader;
    if(!ApiHelp::ReadProcessMemory(hProcess, pMod, &dosHeader, sizeof(dosHeader), &readed))
    {
        return false;
    }

    IMAGE_NT_HEADERS ntHeaders;
    if(!ApiHelp::ReadProcessMemory(hProcess, pMod+dosHeader.e_lfanew, &ntHeaders, sizeof(ntHeaders), &readed))
    {
        return false;
    }

    std::vector<BYTE> imgBuf(ntHeaders.OptionalHeader.SizeOfImage);
    if(!ApiHelp::ReadProcessMemory(hProcess, pMod, &imgBuf[0], imgBuf.size(), &readed))
    {
        return false;
    }

    return BuildModDataFromMemory(&imgBuf[0], modData);
}

//##ModelId=44CF17BE0344
bool ModCheck::BuildModDataFromMemory(BYTE* pMod, ModData& modData)
{
    IMAGE_DOS_HEADER* dosHeader = (IMAGE_DOS_HEADER*)pMod;
    IMAGE_NT_HEADERS* ntHeaders = (IMAGE_NT_HEADERS*)(pMod+dosHeader->e_lfanew);
    IMAGE_SECTION_HEADER* sectHeader = (IMAGE_SECTION_HEADER*)(ntHeaders+1);

    modData.m_headerData.assign((BYTE*)ntHeaders, (BYTE*)(sectHeader+ntHeaders->FileHeader.NumberOfSections));
    
    Crc csum;
    csum.Update(&modData.m_headerData[0], (int)modData.m_headerData.size());
    modData.m_headerHash = csum.Final();

    std::vector<CodeData>::iterator it = modData.m_codeDatas.begin();
    std::vector<CodeData>::iterator e = modData.m_codeDatas.end();
    for(; it != e; ++ it)
    {
        Crc codeSum;
        codeSum.Update(pMod+it->m_rva, it->m_size);
        it->m_hash = codeSum.Final();
    }
    
    return true;
}

} // namespace JoyGuard
